package org.dllwh.milvus.v1;

import cn.hutool.core.util.StrUtil;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DescribeCollectionResponse;
import io.milvus.grpc.GetLoadStateResponse;
import io.milvus.param.R;
import io.milvus.param.RpcStatus;
import io.milvus.param.collection.*;
import io.milvus.param.highlevel.collection.ListCollectionsParam;
import io.milvus.param.highlevel.collection.response.ListCollectionsResponse;
import lombok.extern.slf4j.Slf4j;
import org.dllwh.milvus.model.AjaxResult;
import org.dllwh.milvus.util.MilvusHelper;

import java.util.List;
import java.util.Map;

/**
 * 把今天最好的表现当作明天最新的起点．．～
 * <p>
 * Today the best performance as tomorrow the newest starter!
 *
 * @ClassName MilvusCollectHelper
 * @Description Milvus V1 集合操作
 * @Author: <a href="mailto:duleilewuhen@sina.com">独泪了无痕</a>
 * @Date 2025-05-24 14:33
 * @版本 V 1.0.1
 * @since JDK 17
 */
@Slf4j
public class MilvusCollectHelper {
    /**
     * 判断集合是否已经存在
     *
     * @param collectName 指定集合名称
     * @return 返回值boolean类型, 有(true)/无(false)
     */
    public static boolean hasCollection(String collectName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        HasCollectionParam.Builder requestParam = HasCollectionParam.newBuilder().withCollectionName(collectName);
        R<Boolean> response = milvusClient.hasCollection(requestParam.build());
        if (R.Status.Success.getCode() == response.getStatus()) {
            return response.getData();
        }
        return false;
    }

    /**
     * 创建集合
     */
    public static AjaxResult createCollect(String collectName, List<FieldType> collectFieldList, String description) {
        log.info("Miluvs create collectionName:{}, desc:{}", collectName, description);
        boolean hasCollection = hasCollection(collectName);
        if (hasCollection) {
            return AjaxResult.fail(401, "不能创建重复的结合");
        }
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        CreateCollectionParam.Builder requestParam = CreateCollectionParam.newBuilder().withCollectionName(collectName);

        if (StrUtil.isNotBlank(description)) {
            requestParam.withDescription(description);
        }

        CollectionSchemaParam.Builder collectionSchemaParam = CollectionSchemaParam.newBuilder();
        requestParam.withSchema(collectionSchemaParam.withFieldTypes(collectFieldList).build());

        R<RpcStatus> response = milvusClient.createCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("创建当前连接的数据库中的集合失败，失败原因是：{}", response.getMessage());
        }
        return MilvusHelper.handleRpcStatus(collectName, "创建集合", response);
    }

    /**
     * 获取所有现有Collections
     */
    public static void listCollects() {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        ListCollectionsParam requestParam = ListCollectionsParam.newBuilder().build();
        R<ListCollectionsResponse> response = milvusClient.listCollections(requestParam);
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("获取当前连接的数据库中所有集合失败，失败原因是:{}", response.getMessage());
        }
        if (R.Status.Success.getCode() == response.getStatus()) {

        }
    }

    /**
     * 获取特定集合的详细信息
     *
     * @param collectName 指定集合名称
     */
    public static void describeCollect(String collectName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        DescribeCollectionParam.Builder requestParam = DescribeCollectionParam.newBuilder().withCollectionName(collectName);
        R<DescribeCollectionResponse> response = milvusClient.describeCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("获取特定集合的详细信息失败，失败原因是：{}", response.getMessage());
        }
    }

    /**
     * 集合重新命名
     *
     * @param collectName       指定集合名称
     * @param oldCollectionName 原集合名
     * @param newCollectionName 新集合名
     */
    public static AjaxResult renameCollection(String collectName, String oldCollectionName, String newCollectionName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        RenameCollectionParam.Builder requestParam = RenameCollectionParam.newBuilder()
                .withOldCollectionName(oldCollectionName)
                .withNewCollectionName(newCollectionName);
        R<RpcStatus> response = milvusClient.renameCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("集合重命名失败，失败原因是：{}", response.getMessage());
        }
        return MilvusHelper.handleRpcStatus(collectName, "重命名", response);
    }

    /**
     * 加载集合到内存。必须首先加载集合到内存，否则查询报错。
     *
     * @param collectName 指定集合名称
     */
    public static AjaxResult loadCollect(String collectName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        LoadCollectionParam.Builder requestParam = LoadCollectionParam.newBuilder().withCollectionName(collectName);
        R<RpcStatus> response = milvusClient.loadCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("加载当前连接的数据库中的集合到内存失败，失败原因是:{}", response.getMessage());
        }
        return MilvusHelper.handleRpcStatus(collectName, "集合加载", response);
    }

    /**
     * 从内存中释放集合
     *
     * @param collectName 指定集合名称
     */
    public static AjaxResult releaseCollect(String collectName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        ReleaseCollectionParam.Builder requestParam = ReleaseCollectionParam.newBuilder().withCollectionName(collectName);
        R<RpcStatus> response = milvusClient.releaseCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("释放当前连接的数据库中的集合{}失败，失败原因是:{}", collectName, response.getMessage());
        }
        return MilvusHelper.handleRpcStatus(collectName, "释放", response);
    }

    /**
     * 删除集合
     *
     * @param collectName 指定集合名称
     */
    public static AjaxResult dropCollect(String collectName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        DropCollectionParam.Builder requestParam = DropCollectionParam.newBuilder().withCollectionName(collectName);
        R<RpcStatus> response = milvusClient.dropCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("删除当前连接的数据库中的集合失败，失败原因是：{}", response.getMessage());
        }
        return MilvusHelper.handleRpcStatus(collectName, "删除", response);
    }

    /**
     * 获取当前连接的数据库中集合状态
     *
     * @param collectName 指定集合名称
     */
    public static void getLoadState(String collectName) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        GetLoadStateParam.Builder requestParam = GetLoadStateParam.newBuilder().withCollectionName(collectName);
        R<GetLoadStateResponse> response = milvusClient.getLoadState(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("获取当前连接的数据库中集合状态失败，失败原因是：{}", response.getMessage());
        }
    }

    /**
     * 设置集合属性
     *
     * @param collectName 指定集合名称
     * @param properties  集合属性
     */
    public static AjaxResult alterCollection(String collectName, Map<String, String> properties) {
        MilvusServiceClient milvusClient = MilvusConnectHelper.getMilvusServiceClient();
        AlterCollectionParam.Builder requestParam = AlterCollectionParam.newBuilder().withCollectionName(collectName);
        properties.forEach(requestParam::withProperty);
        R<RpcStatus> response = milvusClient.alterCollection(requestParam.build());
        if (R.Status.Success.getCode() != response.getStatus()) {
            log.error("设置集合属性失败，失败原因是：{}", response.getMessage());
        }
        return MilvusHelper.handleRpcStatus(collectName, "更新", response);
    }
}
